import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from copy import deepcopy

# 定义网格大小和参数
grid_size = 10
gamma = 0.9
theta = 1e-4

# 初始化效用值
U = np.zeros((grid_size, grid_size))

# 终止状态
goal_state = (9, 9)  # 注意索引从0开始，(10,10)对应(9,9)
U[goal_state] = 1

# 值迭代算法
def value_iteration(U):
    delta = float('inf')
    while delta > theta:
        delta = 0
        U_old = deepcopy(U)
        for x in range(grid_size):
            for y in range(grid_size):
                if (x, y) == goal_state:
                    continue
                v = U_old[x, y]
                # 计算四个方向的效用值
                utility_values = []
                for action in [(-1,0), (1,0), (0,-1), (0,1)]:
                    x_new, y_new = x + action[0], y + action[1]
                    if 0 <= x_new < grid_size and 0 <= y_new < grid_size:
                        utility = U_old[x_new, y_new]
                    else:
                        utility = U_old[x, y]
                    utility_values.append(utility)
                U[x, y] = gamma * max(utility_values)
                delta = max(delta, abs(v - U[x, y]))
    return U

U = value_iteration(U)

# 构建特征矩阵 X
X = []
Y = []
for x in range(grid_size):
    for y in range(grid_size):
        X.append([1, x+1, y+1])
        Y.append(U[x, y])
X = np.array(X)
Y = np.array(Y)

# 最小二乘法求解 w
w = np.linalg.inv(X.T @ X) @ X.T @ Y

# 计算线性逼近的效用值
U_hat = X @ w
U_hat_grid = U_hat.reshape((grid_size, grid_size))

# 绘制三维图形
fig = plt.figure(figsize=(12, 6))

# 真实效用函数
ax1 = fig.add_subplot(121, projection='3d')
X_plot, Y_plot = np.meshgrid(range(1, grid_size+1), range(1, grid_size+1))
ax1.plot_surface(X_plot, Y_plot, U, cmap='viridis')
ax1.set_title('真实效用函数')

# 线性逼近效用函数
ax2 = fig.add_subplot(122, projection='3d')
ax2.plot_surface(X_plot, Y_plot, U_hat_grid, cmap='viridis')
ax2.set_title('线性逼近效用函数')

plt.show()